package com.yf.service.chat.impl;

import cn.hutool.json.JSONUtil;
import com.yf.model.ai.entity.AiConversation;
import com.yf.model.ai.entity.AiMask;
import com.yf.model.ai.entity.AiMessage;
import com.yf.model.ai.enums.AiMessageRoleEnum;
import com.yf.service.manage.IAiConversationService;
import com.yf.service.manage.IAiMaskService;
import com.yf.service.manage.IAiMessageService;
import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.jetbrains.annotations.NotNull;
import org.springframework.ai.chat.memory.ChatMemory;
import org.springframework.ai.chat.messages.AssistantMessage;
import org.springframework.ai.chat.messages.Message;
import org.springframework.ai.chat.messages.SystemMessage;
import org.springframework.ai.chat.messages.UserMessage;
import org.springframework.stereotype.Component;
import org.springframework.transaction.annotation.Transactional;

import java.time.LocalDateTime;
import java.util.ArrayList;
import java.util.List;

/**
 * Ai Message Table Chat Memory
 *
 * @author : YiFei
 * @since : 2025/6/9 23:17
 */
@Slf4j
@Component
@RequiredArgsConstructor
public class AiMessageChatMemory implements ChatMemory {

    private final IAiConversationService conversationService;
    private final IAiMessageService messageService;
    private final IAiMaskService maskService;

    @Override
    public void add(@NotNull String conversationId, List<Message> messages) {
        // 检查会话是否存在
        for (Message message : messages) {
            // 创建新的消息实体
            AiMessage aiMessage = new AiMessage();
            aiMessage.setConversationId(Long.valueOf(conversationId));
            // 设置消息角色
            if (message instanceof UserMessage) {
                aiMessage.setRole(AiMessageRoleEnum.USER);
                // 更新会话活跃时间
                conversationService.lambdaUpdate()
                        .set(AiConversation::getLastActiveTime, LocalDateTime.now())
                        .eq(AiConversation::getId, conversationId).update();
            } else if (message instanceof AssistantMessage) {
                aiMessage.setRole(AiMessageRoleEnum.ASSISTANT);
            } else if (message instanceof SystemMessage) {
                aiMessage.setRole(AiMessageRoleEnum.SYSTEM);
            } else {
                log.warn("Unsupported message type: {}", message.getClass().getSimpleName());
                continue; // 跳过不支持的消息类型
            }
            // 设置消息内容
            aiMessage.setContent(message.getText());
            // 设置 Metadata
            aiMessage.setMetadata(JSONUtil.toJsonStr(message.getMetadata()));
            // 设置发送时间
            aiMessage.setSendTime(LocalDateTime.now());
            // 保存消息到数据库
            messageService.save(aiMessage);
        }
        if (messageService.lambdaQuery().eq(AiMessage::getConversationId,conversationId).count() == 1) {
            // 如果是第一次添加消息，更新会话名为消息的前10个字符
            Message message = messages.get(0);
            conversationService.lambdaUpdate()
                    .eq(AiConversation::getId, conversationId)
                    .set(AiConversation::getTitle,message.getText().substring(0, Math.min(message.getText().length(), 10)))
                    .update();
        }
    }

    @Override
    public List<Message> get(@NotNull String conversationId) {
        AiConversation conversation = conversationService.lambdaQuery()
                .eq(AiConversation::getId, conversationId)
                .one();
        List<Message> messages = new ArrayList<>();
        // 查询会话面具
        Integer currentMaskId = conversation.getCurrentMaskId();
        if (currentMaskId != null) {
            AiMask mask = maskService.getById(currentMaskId);
            if (mask != null) {
                // 如果面具存在，添加面具的消息
                Message message = new SystemMessage(mask.getPrompt());
                messages.add(message);
            } else {
                log.warn("Mask with ID {} not found for conversation {}", currentMaskId, conversationId);
            }
        }

        // 查询过往聊天记录 , 仅查询 AssistantMessage 、UserMessage
        messageService.lambdaQuery()
                .eq(AiMessage::getConversationId, conversationId)
                .in(AiMessage::getRole, AiMessageRoleEnum.ASSISTANT, AiMessageRoleEnum.USER)
                .orderByAsc(AiMessage::getSendTime)
                .list()
                .forEach(aiMessage -> {
                    if (AiMessageRoleEnum.ASSISTANT.equals(aiMessage.getRole())) {
                        // 如果是助手消息，添加到消息列表
                        AssistantMessage assistantMessage = new AssistantMessage(aiMessage.getContent());// 设置消息ID
                        messages.add(assistantMessage);
                    } else if (AiMessageRoleEnum.USER.equals(aiMessage.getRole())) {
                        // 如果是用户消息，添加到消息列表
                        Message message = new UserMessage(aiMessage.getContent());
                        messages.add(message);
                    }
                });


        return messages;
    }

    @Override
    @Transactional
    public void clear(@NotNull String conversationId) {
        // 删除会话信息
        conversationService.removeById(conversationId);
        // 删除聊天记录
        messageService.lambdaUpdate()
                .eq(AiMessage::getConversationId, conversationId)
                .remove();
    }
}
